#!/usr/bin/python

# ditherer.py
# prototype script for dithering of images...
# uses different working areas for L and A&B channels
# created by tiborb95@gmail.com
# GPL v.3
# version 0.1
# it is work in progress...
# needs the CS_convertor.py file in the same directory
# should be compatible with python 2.7 and above



######### importing and setting global variables ###########
from __future__ import division  # to return float numbers after division allways
from sys import argv,exit,stdout
from getopt import getopt
from os import path,getenv
from os.path import split,splitext
import numpy
import CS_convertor

version="0.1.0"

Lcur=numpy.array([	[ 0, 1, 1, 1],
					[ 0, 2, 3, 1],
					[ 1, 3, 10, 0]	])

abcur=numpy.array([	[0, 0, 1, 1, 1, 0 ],
					[0, 1, 2, 2, 2, 1 ],
					[1, 2, 3, 4, 3, 1 ],
					[1, 2, 4, 4, 0, 0 ]])

Lspread=numpy.array([	[ 0, 0, 5, 1],
						[ 1, 5, 4, 1],
						[ 0, 1, 1, 0]])

ABspread=numpy.array([	[ 0, 0, 4, 2, 1],
						[ 2, 4, 3, 1, 1],
						[ 1, 2, 2, 1, 0],
						[ 0, 1, 1, 0, 0]])

Lspreadsum=numpy.sum(Lspread)
ABspreadsum=numpy.sum(ABspread)
quiet=True
verbose=False
halfmode=False
level=3
lweight=1  #l channel weight



try:
	from PIL import Image      # necessary for manipulation with images
	from PIL import ImageFilter
	from PIL import ImageDraw,ImageFont
except:
	print (" Warning: Failed to load Pillow (Imaging) modules!")
	print (" Pillow (fork of PIL) should be available in repositories")
	print (" of your distribution, please install it first.")
	exit()

newsize=0    #no resize

def spread_error(error,array,base,mask,x,y,xrel,yrel,verbose):
	if not (verbose):
		print (" Spreading error: "+str(error)+", divider: "+str(base))
		print ("Doing: "+str(x)+"x"+str(y))
	for fcol in range(len(mask)):
		for frow in range(len(mask[fcol])):	
			localx=x+xrel+fcol
			localy=y+yrel+frow
			if (localx<0 or localy<0 or localx >= imgX or localy>=imgY):
				continue	
			if not (verbose):
				print ("  Weight: "+str(mask[fcol,frow]))
				print ("  Modifying: "+str(x+xrel+fcol)+"x"+str(y+yrel+frow))
				print (	"   Curr error:" + str(array[localx,localy]))
			#print (mask[fcol,frow])
			array[localx,localy]=array[localx,localy]+float(error*mask[fcol,frow]/base)
			if not verbose:
				print (	"   New error:" + str(array[localx,localy]))
			
			
			
			
def get_sum(array,mask,x,y,xrel,yrel,verbose):
	sum=0
	count=0
	if not verbose:
		print ("Doing: "+str(x)+"x"+str(y))
	for fcol in range(len(mask)):
		for frow in range(len(mask[fcol])):
			localx=x+xrel+fcol
			localy=y+yrel+frow
			if not verbose:
				print (" Scanning: "+str(x+xrel+fcol)+"x"+str(y+yrel+frow))
			if (localx<0 or localy<0 or localx >= imgX or localy>=imgY):
				continue
			sum=sum+array[localx,localy]
			count=count+1
	if not verbose:
		print ("  orig value "+str(array[x,y]))
		print ("  perceived value "+str(sum/count))	
	return 	sum/count
	


print ("[ditherer.py, version: "+str(version)+", use '-h' for help and bit more info]")

files=[]
#parsing arguments
for arg in argv[1:]:
	if arg.startswith("-s"):
		#print (arg[2:])
		try:
			newsize=int(arg[2:])
			print (" New size of images: "+str(newsize))
		except:
			print ("! Failed to get new size:" + str(arg[2:])+" ???")
			newsize=0
		if newsize==0:
			pass
		elif (newsize<10):
			newsize=10
		if (newsize>2500):
			newsize=2500
	if arg.startswith("-l"):
		#print (arg[2:])
		try:
			tmp=int(arg[2:])
			#print (" New size of images: "+str(newsize))
		except:
			print ("! Failed to get color count:" + str(arg[2:])+" ???")
			continue

		if tmp>=2 and tmp<=10:
			pass
		elif (tmp<2):
			tmp=2
		elif (tmp>10):
			tmp=10
		level=tmp
		print (" Using level (color count): "+str(level))
		if level>=6:
			print (" ! Note: processing time grows exponentialy based on level (color count)")
	if arg.startswith("-L"):
		#print (arg[2:])
		try:
			tmp=float(arg[2:])
			#print (" New size of images: "+str(newsize))
		except:
			print ("! Failed to get L channel weight:" + str(arg[2:])+" ???")
			continue

		if tmp>=0.1 and tmp<=10:
			pass
		elif (tmp<0.1):
			tmp=0.1
		elif (tmp>10):
			tmp=10
		lweight=tmp
		print (" L channel weight to be used: "+str(lweight))

	if arg == "-H":
		print (" Halfmode...")
		halfmode=True
	if arg == "-h":
		print (" ")
		print ("ditherer.py "+version)
		print ("Available switches")
		print ("-s#  where # is width of image (no space)")	
		print ("-H  halfmode")	
		#print ("-S	put statistics on the image")
		print ("All switches not starting with '-' are considered images")
		print ("You might force specific version of python by prepending a binary like 'python2.7' before the command itself")
		print (" ")
		print ("This is prototype/alpha level software... ")	
		print ("Contact: tiborb95 at gmail dot com ")	
		exit (0)		
	
	if not arg.startswith("-"):
		files.append(arg)


#generate starting colors, for now it is combination of 0,128,255 RGB values
colors=[]
levels=[]
levels.append(0)
for i in range(1,level-1):
	levels.append(int(round(i/(level-1)*255)))
levels.append(255)
print (levels)



levels
for R in levels:
	for G in levels:
		for B in levels:
			l,a,b = CS_convertor.rgb_to_cielab(R,G,B)
			colors.append([l,a,b])

		

if len(files)==0:
	print ("No files/images given...")


imgcount=0			
for image in files:
	imgcount=imgcount+1	

	corename=splitext(split(image)[1])[0]

	# trying to open image
	try:
		sourceImg = Image.open(image)
		sourceImg_loaded=sourceImg.load()
		print ("\nProcessing: "+image+"      "+str(imgcount)+"/"+str(len(files)))
	except:
		print ("Cannot open "+image+" image, skipping this one...")
		continue
	
	
	#resizing if needed
	imgX,imgY=sourceImg.size
	if newsize>0 and imgX > newsize:
		print (" Resizing...("+str(newsize)+")")
		sourceImg=sourceImg.resize((newsize,int(newsize/imgX*imgY)),Image.ANTIALIAS)
		sourceImg_loaded=sourceImg.load()
	
	
	# BODY
	imgX,imgY=sourceImg.size
	estat=[0,0,0] #errors for LAB channels

	# creating new image
	newImage = Image.new( 'RGB', (imgX,imgY), "black") #


	sourceL=numpy.zeros((imgX,imgY))
	sourceA=numpy.zeros((imgX,imgY))
	sourceB=numpy.zeros((imgX,imgY))
	finalLAB=numpy.zeros((imgX,imgY,3))
	finalRGB=numpy.zeros((imgX,imgY,3))
	Lerrors=numpy.zeros((imgX,imgY))
	Aerrors=numpy.zeros((imgX,imgY))	
	Berrors=numpy.zeros((imgX,imgY))	
		
	for x in range(imgX):
		for y in range(imgY):
			R=sourceImg_loaded[x,y][0]
			G=sourceImg_loaded[x,y][1]
			B=sourceImg_loaded[x,y][2]
			sourceL[x,y],sourceA[x,y],sourceB[x,y] = CS_convertor.rgb_to_cielab(R,G,B)

	Lorigtotsum=0
	Aorigtotsum=0
	Borigtotsum=0
	Lnewtotsum=0
	Anewtotsum=0
	Bnewtotsum=0	

	
	#iterating over image and calculating current (perceived) values
	for y in range(imgY):
	
		if y%200==0:
			print ("  "+str(y)+"/"+str(imgY))

		Lorigsum=0
		Aorigsum=0
		Borigsum=0
		Lnewsum=0
		Anewsum=0
		Bnewsum=0		
		
		for x in range(imgX):
			Lv=get_sum(sourceL,Lcur,x,y,-2,-2,quiet)
			Av=get_sum(sourceA,abcur,x,y,-3,-3,quiet)
			Bv=get_sum(sourceB,abcur,x,y,-3,-3,quiet)
			
			Lorigsum=Lorigsum+sourceL[x,y]
			Aorigsum=Aorigsum+sourceA[x,y]
			Borigsum=Borigsum+sourceB[x,y]
			Lorigtotsum=Lorigtotsum+sourceL[x,y]
			Aorigtotsum=Aorigtotsum+sourceA[x,y]
			Borigtotsum=Borigtotsum+sourceB[x,y]
			Lv=Lv+Lerrors[x,y]
			Av=Av+Aerrors[x,y]
			Bv=Bv+Berrors[x,y]
			#print (Lv,Av,Bv)
				
							
			#finding the best color
			mindiff=1000
			bestcol=[0,0,0]
			for color in colors:
				diff=lweight*abs(color[0]-Lv) + abs(color[1]-Av) + abs(color[2]-Bv)
				if diff<mindiff:
					mindiff=diff
					bestcol=color

			finalLAB[x,y]=bestcol
			
			Lnewsum=Lnewsum+finalLAB[x,y,0]
			Anewsum=Anewsum+finalLAB[x,y,1]
			Bnewsum=Bnewsum+finalLAB[x,y,2]
			Lnewtotsum=Lnewtotsum+finalLAB[x,y,0]
			Anewtotsum=Anewtotsum+finalLAB[x,y,1]
			Bnewtotsum=Bnewtotsum+finalLAB[x,y,2]
			
			Lerror=Lv-bestcol[0]
			Aerror=Av-bestcol[1]			
			Berror=Bv-bestcol[2]

			estat[0]=estat[0]+abs(Lerror)
			estat[1]=estat[1]+abs(Aerror)
			estat[2]=estat[2]+abs(Berror)
			
			#saving errors 
			spread_error(Lerror,Lerrors,Lspreadsum,Lspread,x,y,-1,0,quiet)
			spread_error(Aerror,Aerrors,ABspreadsum,ABspread,x,y,-1,0,quiet)			
			spread_error(Berror,Berrors,ABspreadsum,ABspread,x,y,-1,0,quiet)
		
		#statistics for row
		#print ("  Stat: orig: "+str(round(Lorigsum/imgX)) + " - "+str(round(Aorigsum/imgX))+ " - "+str(round(Borigsum/imgX)))
		#print ("         new: "+str(round(Lnewsum /imgX)) + " - "+str(round( Anewsum/imgX))+ " - "+str(round( Bnewsum/imgX)))
	#converting LAB to RGB, and saving results into newImage
	print ("  Converting to RGB and saving...")
	for x in range(imgX):
		for y in range(imgY):
			if halfmode==True and x<imgX/2.5:
				R=sourceImg_loaded[x,y][0]
				G=sourceImg_loaded[x,y][1]
				B=sourceImg_loaded[x,y][2]
				newImage.putpixel((x,y),(R,G,B))
				continue
			R,G,B=CS_convertor.cielab_to_rgb(finalLAB[x,y,0],finalLAB[x,y,1],finalLAB[x,y,2])
			newImage.putpixel((x,y),(int(R) , int(G) , int(B)))



	#saving
	newname="dith_"+corename+"_C"+str(len(colors))+".jpg"
	print( "  Saving as: "+newname)
	newImage.save(newname,quality=95)
	newname="dith_"+corename+"_C"+str(len(colors))+".png"
	print( "  Saving as: "+newname)
	newImage.save(newname,optimize=True)
	
	#printing statistics
	print (" SUMMARY for image         L    A      B")
	print ("  Average error       : "+str(round(estat[0]/imgX/imgY,1))+"   "+str(round(estat[1]/imgX/imgY,1))+"   "+ str(round(estat[2]/imgX/imgY,1))+"   (average differences)")	
	print ("  Averages source img : "+str(round(Lorigtotsum/imgX/imgY,1)) + "   "+str(round(Aorigtotsum/imgX/imgY))+ "   "+str(round(Borigtotsum/imgX/imgY,1))+"   (average values)")
	print ("           final img  : "+str(round(Lnewtotsum /imgX/imgY,1)) + "   "+str(round( Anewtotsum/imgX/imgY))+ "   "+str(round( Bnewtotsum/imgX/imgY,1)))
